#include "refine/despill.h"
#include "refine/window_pass.h"
#include <limits>
//#include <iostream>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <glib/g2/smatrix.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace despill {

typedef unsigned uint;
typedef long long llong;

const int uchar_max = std::numeric_limits<unsigned char>::max();

inline unsigned char uchar_limit(float c)
{
	return (c + 0.5f >= uchar_max) ? uchar_max : (c <= 0) ? 0 : c + 0.5f;
}

// reconstructs foreground color from a number of image samples with known alpha
template <typename Raster>
class despiller
{
public:
	typedef typename Raster::elem_type color;
	typedef typename color::component_type color_component;
public:
	despiller(Raster & im_, params const & p_) : im(im_), prm(p_) {
		clear(); 

		img::bitmap not_max(im.dim());
		for (img::bitmap::iterator i = not_max.begin(); i != not_max.end(); ++i)
			*i = im[i].a != uchar_max;

		img::bitmap not_min(im.dim());
		for (img::bitmap::iterator i = not_min.begin(); i != not_min.end(); ++i)
			*i = im[i].a != 0;
		uint w = 1 + prm.whs * (prm.max_condition_number - 1) / (prm.max_condition_number + 1);
		img::dilate4(not_min, w);

		modifiable = not_max;
		img::combine(modifiable, not_min, img::bit_and());
	}
	void clear() {
		n = sum_alpha = sum_sqr_alpha = 0;
		sum_I = sum_alpha_I = img::BGR<uint>(0, 0, 0);
	}
	size_extent dim() const
		{ return im.dim(); }
	bool selected(size_point const & p) const
		{ return modifiable[p]; /*im[p].a != uchar_max && im[p].a != 0;*/ }
	void add(size_point const & p)
		{ add_sample(im[p]); }
	void remove(size_point const & p)
		{ remove_sample(im[p]); }
	void remove_all() 
		{ clear(); }
	void process(size_point const & p, size_rect const &) 
	{
		//E = A^T * A
		g2::smatrixf E(
			sum_sqr_alpha, //xx
			uchar_max * sum_alpha - sum_sqr_alpha, //xy
			uchar_max * uchar_max * n - 2 * uchar_max * sum_alpha + sum_sqr_alpha); //yy

		color & I = im[p];

		if (!prm.mimicry) {
			float D = sqr(E.xx - E.yy) + 4 * E.xy * E.xy;
			if (sqr(prm.max_condition_number + 1) * D > sqr((prm.max_condition_number - 1) * E.tr()))
				return;

			//top row (foreground) of inversed E
			g2::vectorf m = float(uchar_max) / E.det() * g2::vectorf(E.yy, -E.xy);

			g2::vectorf f;

			// blue
			f = g2::vectorf(sum_alpha_I.b, uchar_max * sum_I.b - sum_alpha_I.b);
			I.b = uchar_limit(scalar_product(m, f));

			// green
			f = g2::vectorf(sum_alpha_I.g, uchar_max * sum_I.g - sum_alpha_I.g);
			I.g = uchar_limit(scalar_product(m, f));

			// red
			f = g2::vectorf(sum_alpha_I.r, uchar_max * sum_I.r - sum_alpha_I.r);
			I.r = uchar_limit(scalar_product(m, f));

			return;
		}

		uint a0 = I.a;

		g2::vectorf C(a0, uchar_max - a0);
		g2::smatrixf CC = float(n) * tsquare(C);
		float cc = 1 / CC.tr(); //(C * C^T)^-1

		//S = I - C^T * (C * C^T)^-1 * C
		g2::smatrixf S(CC.yy, -CC.xy, CC.xx);
		S *= cc;

		g2::smatrixf M = CC + duplex_product(E, S); // M = C^T * C + S * E * S

		float D = sqr(M.xx - M.yy) + 4 * M.xy * M.xy;
		if (sqr(prm.max_condition_number + 1) * D > sqr((prm.max_condition_number - 1) * M.tr()))
			return;

		//top row (foreground) of inversed M
		g2::vectorf m = float(uchar_max) / M.det() * g2::vectorf(M.yy, -M.xy);
		g2::vectorf mS = S * m;
		float mEc = float(n) * scalar_product(m, (C - cc * (S * (E * C)))); // = (I - (C * C^T)^-1 * S * E) C^T

		g2::vectorf f;

		// blue
		f = g2::vectorf(sum_alpha_I.b, uchar_max * sum_I.b - sum_alpha_I.b);
		I.b = uchar_limit(scalar_product(mS, f) + float(I.b) * mEc);

		// green
		f = g2::vectorf(sum_alpha_I.g, uchar_max * sum_I.g - sum_alpha_I.g);
		I.g = uchar_limit(scalar_product(mS, f) + float(I.g) * mEc);

		// red
		f = g2::vectorf(sum_alpha_I.r, uchar_max * sum_I.r - sum_alpha_I.r);
		I.r = uchar_limit(scalar_product(mS, f) + float(I.r) * mEc);
	}
private:
	void add_sample(color const & I) {
		++n;
		uint alpha = I.a;
		sum_alpha += alpha;
		sum_sqr_alpha += alpha * alpha;
		sum_I += I;
		sum_alpha_I += alpha * img::BGR<uint>(I);
	}
	void remove_sample(color const & I) {
		--n;
		uint alpha = I.a;
		sum_alpha -= alpha;
		sum_sqr_alpha -= alpha * alpha;
		sum_I -= I;
		sum_alpha_I -= alpha * img::BGR<uint>(I);
	}
	bool singular() const
		{ return det() == 0; }
	llong det() const
		{ return llong(n) * llong(sum_sqr_alpha) - llong(sum_alpha) * llong(sum_alpha); }
private:
	params prm;
	Raster & im;
	uint n, sum_alpha, sum_sqr_alpha;
	img::BGR<uint> sum_I, sum_alpha_I;
	img::bitmap modifiable; ///< this pixels may change their alpha theoretically
};

void perform(byte_image & im, params const & p)
{
	despiller<byte_image> ds(im, p);
	window_pass(ds, p.whs);
}

void perform(argb_byte_image & im, params const & p)
{
	despiller<argb_byte_image> ds(im, p);
	window_pass(ds, p.whs);
}

} //namespace despill

} //namespace refine
